/*******************************************************
 * Copyright (C) 2019, Aerial Robotics Group, Hong Kong University of Science
 *and Technology
 *
 * This file is part of VINS.
 *
 * Licensed under the GNU General Public License v3.0;
 * you may not use this file except in compliance with the License.
 *******************************************************/

#pragma once
#include <ros/assert.h>
#include <eigen3/Eigen/Dense>
#include <iostream>

#include "lci_slam/vilens/ceres/IntegrationBase.h"
#include "lci_slam/vilens/utility/utility.h"

#include <ceres/ceres.h>

class ImuFactor : public ceres::SizedCostFunction<15, 7, 9, 7, 9> {
public:
    enum StateOrder { O_P = 0, O_R = 3, O_V = 6, O_BA = 9, O_BG = 12 };
    enum NoiseOrder { O_AN = 0, O_GN = 3, O_AW = 6, O_GW = 9 };
    Eigen::Vector3d G = Eigen::Vector3d(0, 0, 9.81);

public:
    ImuFactor() = delete;
    ImuFactor(const IntegrationBase::Ptr &_pre_integration)
        : pre_integration(_pre_integration) {}
    virtual bool Evaluate(double const *const *parameters, double *residuals,
                          double **jacobians) const {
        Eigen::Vector3d Pi(parameters[0][0], parameters[0][1],
                           parameters[0][2]);
        Eigen::Quaterniond Qi(parameters[0][6], parameters[0][3],
                              parameters[0][4], parameters[0][5]);

        Eigen::Vector3d Vi(parameters[1][0], parameters[1][1],
                           parameters[1][2]);
        Eigen::Vector3d Bai(parameters[1][3], parameters[1][4],
                            parameters[1][5]);
        Eigen::Vector3d Bgi(parameters[1][6], parameters[1][7],
                            parameters[1][8]);

        Eigen::Vector3d Pj(parameters[2][0], parameters[2][1],
                           parameters[2][2]);
        Eigen::Quaterniond Qj(parameters[2][6], parameters[2][3],
                              parameters[2][4], parameters[2][5]);

        Eigen::Vector3d Vj(parameters[3][0], parameters[3][1],
                           parameters[3][2]);
        Eigen::Vector3d Baj(parameters[3][3], parameters[3][4],
                            parameters[3][5]);
        Eigen::Vector3d Bgj(parameters[3][6], parameters[3][7],
                            parameters[3][8]);

        Eigen::Map<Eigen::Matrix<double, 15, 1>> residual(residuals);
        residual = pre_integration->evaluate(Pi, Qi, Vi, Bai, Bgi, Pj, Qj, Vj,
                                             Baj, Bgj);

        Eigen::Matrix<double, 15, 15> sqrt_info =
            Eigen::LLT<Eigen::Matrix<double, 15, 15>>(
                pre_integration->covariance.inverse())
                .matrixL()
                .transpose();
        residual = sqrt_info * residual;

        if (jacobians) {
            double sum_dt = pre_integration->sum_dt;
            Eigen::Matrix3d dp_dba =
                pre_integration->jacobian.template block<3, 3>(O_P, O_BA);
            Eigen::Matrix3d dp_dbg =
                pre_integration->jacobian.template block<3, 3>(O_P, O_BG);

            Eigen::Matrix3d dq_dbg =
                pre_integration->jacobian.template block<3, 3>(O_R, O_BG);

            Eigen::Matrix3d dv_dba =
                pre_integration->jacobian.template block<3, 3>(O_V, O_BA);
            Eigen::Matrix3d dv_dbg =
                pre_integration->jacobian.template block<3, 3>(O_V, O_BG);

            if (pre_integration->jacobian.maxCoeff() > 1e8 ||
                pre_integration->jacobian.minCoeff() < -1e8) {
                ROS_WARN("numerical unstable in preintegration");
            }

            if (jacobians[0]) {
                Eigen::Map<Eigen::Matrix<double, 15, 7, Eigen::RowMajor>>
                    jacobian_pose_i(jacobians[0]);
                jacobian_pose_i.setZero();

                jacobian_pose_i.block<3, 3>(O_P, O_P) =
                    -Qi.inverse().toRotationMatrix();
                jacobian_pose_i.block<3, 3>(O_P, O_R) = Utility::skewSymmetric(
                    Qi.inverse() *
                    (0.5 * G * sum_dt * sum_dt + Pj - Pi - Vi * sum_dt));

                Eigen::Quaterniond corrected_delta_q =
                    pre_integration->delta_q *
                    Utility::deltaQ(dq_dbg *
                                    (Bgi - pre_integration->linearized_bg));
                jacobian_pose_i.block<3, 3>(O_R, O_R) =
                    -(Utility::Qleft(Qj.inverse() * Qi) *
                      Utility::Qright(corrected_delta_q))
                         .bottomRightCorner<3, 3>();
                jacobian_pose_i.block<3, 3>(O_V, O_R) = Utility::skewSymmetric(
                    Qi.inverse() * (G * sum_dt + Vj - Vi));

                jacobian_pose_i = sqrt_info * jacobian_pose_i;

                if (jacobian_pose_i.maxCoeff() > 1e8 ||
                    jacobian_pose_i.minCoeff() < -1e8) {
                    ROS_WARN("numerical unstable in preintegration");
                }
            }
            if (jacobians[1]) {
                Eigen::Map<Eigen::Matrix<double, 15, 9, Eigen::RowMajor>>
                    jacobian_speedbias_i(jacobians[1]);
                jacobian_speedbias_i.setZero();
                jacobian_speedbias_i.block<3, 3>(O_P, O_V - O_V) =
                    -Qi.inverse().toRotationMatrix() * sum_dt;
                jacobian_speedbias_i.block<3, 3>(O_P, O_BA - O_V) = -dp_dba;
                jacobian_speedbias_i.block<3, 3>(O_P, O_BG - O_V) = -dp_dbg;

                jacobian_speedbias_i.block<3, 3>(O_R, O_BG - O_V) =
                    -Utility::Qleft(Qj.inverse() * Qi *
                                    pre_integration->delta_q)
                         .bottomRightCorner<3, 3>() *
                    dq_dbg;

                jacobian_speedbias_i.block<3, 3>(O_V, O_V - O_V) =
                    -Qi.inverse().toRotationMatrix();
                jacobian_speedbias_i.block<3, 3>(O_V, O_BA - O_V) = -dv_dba;
                jacobian_speedbias_i.block<3, 3>(O_V, O_BG - O_V) = -dv_dbg;

                jacobian_speedbias_i.block<3, 3>(O_BA, O_BA - O_V) =
                    -Eigen::Matrix3d::Identity();

                jacobian_speedbias_i.block<3, 3>(O_BG, O_BG - O_V) =
                    -Eigen::Matrix3d::Identity();

                jacobian_speedbias_i = sqrt_info * jacobian_speedbias_i;
            }
            if (jacobians[2]) {
                Eigen::Map<Eigen::Matrix<double, 15, 7, Eigen::RowMajor>>
                    jacobian_pose_j(jacobians[2]);
                jacobian_pose_j.setZero();

                jacobian_pose_j.block<3, 3>(O_P, O_P) =
                    Qi.inverse().toRotationMatrix();

                Eigen::Quaterniond corrected_delta_q =
                    pre_integration->delta_q *
                    Utility::deltaQ(dq_dbg *
                                    (Bgi - pre_integration->linearized_bg));
                jacobian_pose_j.block<3, 3>(O_R, O_R) =
                    Utility::Qleft(corrected_delta_q.inverse() * Qi.inverse() *
                                   Qj)
                        .bottomRightCorner<3, 3>();

                jacobian_pose_j = sqrt_info * jacobian_pose_j;
            }
            if (jacobians[3]) {
                Eigen::Map<Eigen::Matrix<double, 15, 9, Eigen::RowMajor>>
                    jacobian_speedbias_j(jacobians[3]);
                jacobian_speedbias_j.setZero();

                jacobian_speedbias_j.block<3, 3>(O_V, O_V - O_V) =
                    Qi.inverse().toRotationMatrix();

                jacobian_speedbias_j.block<3, 3>(O_BA, O_BA - O_V) =
                    Eigen::Matrix3d::Identity();

                jacobian_speedbias_j.block<3, 3>(O_BG, O_BG - O_V) =
                    Eigen::Matrix3d::Identity();

                jacobian_speedbias_j = sqrt_info * jacobian_speedbias_j;
            }
        }

        return true;
    }
    IntegrationBase::Ptr pre_integration;
};
